import cv2, os, requests, io, math, operator
import numpy as np
from PIL import Image
from support.utils.dir import *
from functools import reduce

class StarConfidence:
    def __init__(self):
        self.starResList = self.__loadRes(os.path.dirname(os.path.abspath(__file__)))

    def __starnum(self, elem):
        return elem.split('.')[0]

    def __loadRes(self, path):
        resList = []
        filepath = os.path.join(path, 'res')
        for file in sorted(os.listdir(filepath), key=self.__starnum):
            resList.append(Image.open(os.path.join(filepath,file)).histogram())
            #resList.append(cv2.imread(os.path.join(filepath, file)))
        return resList

    def clean(self):
        for array in self.starResList:
            del array

    def getStarNum(self, url):
        res = requests.get(url)
        imgBuf = io.BytesIO(res.content)
        histogram = Image.open(imgBuf).histogram()
        num = 0
        for i in range(len(self.starResList)):
            #confidence = self.__template_match_in(tmp, self.starResList[i])
            confidence = math.sqrt(reduce(operator.add, list(map(lambda a,b: (a-b)**2, self.starResList[i], histogram)))/len(histogram))
            if(confidence == 0):
                num = i
                break
        return num

    def __template_match_in(self, templ, img_src):
        # if(Global.MATCH_THROUGH_GRAY):
            # img_src = cv2.cvtColor(img_src, cv2.COLOR_BGR2GRAY)
            # templ = cv2.cvtColor(templ, cv2.COLOR_BGR2GRAY)
        confidence = 0
        h, w = templ.shape[:2]
        for match_method in [cv2.TM_CCOEFF, cv2.TM_CCOEFF_NORMED, #cv2.TM_CCORR,
            cv2.TM_CCORR_NORMED, cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]:
            result = cv2.matchTemplate(img_src, templ, match_method)
            cv2.normalize( result, result, 0, 1, cv2.NORM_MINMAX, -1 )
            _minVal, _maxVal, minLoc, maxLoc = cv2.minMaxLoc(result, None)
            if (match_method == cv2.TM_SQDIFF or match_method == cv2.TM_SQDIFF_NORMED):
                matchLoc = minLoc
            else:
                matchLoc = maxLoc
            # 求取可信度:
            _confidence = self.__get_confidence_from_matrix(img_src, templ, matchLoc, _maxVal, w, h)
            # print("confidence:{}, pos:{} size:{}x{}".format(_confidence, matchLoc, matchLoc[0] + templ.shape[1], matchLoc[1] + templ.shape[0]))
            if(_confidence > confidence):
                confidence = _confidence
        return confidence
    
    def __get_confidence_from_matrix(self, im_source, im_search, max_loc, max_val, w, h):
        """根据结果矩阵求出confidence."""
        # 求取可信度:
        # if Global.MATCH_THROUGH_GRAY:
            # confidence = max_val
        # else:
            # 如果有颜色校验,对目标区域进行BGR三通道校验:
        img_crop = im_source[max_loc[1]:max_loc[1] + h, max_loc[0]: max_loc[0] + w]
        confidence = self.__cal_rgb_confidence(img_crop, im_search)
        return confidence
    
    def __cal_rgb_confidence(self, img_src_rgb, img_sch_rgb):
        """同大小彩图计算相似度."""
        # 扩展置信度计算区域
        img_sch_rgb = cv2.copyMakeBorder(img_sch_rgb, 10,10,10,10,cv2.BORDER_REPLICATE)
        # 转HSV强化颜色的影响
        img_src_rgb = cv2.cvtColor(img_src_rgb, cv2.COLOR_BGR2HSV)
        img_sch_rgb = cv2.cvtColor(img_sch_rgb, cv2.COLOR_BGR2HSV)
        src_bgr, sch_bgr = cv2.split(img_src_rgb), cv2.split(img_sch_rgb)

        # 计算BGR三通道的confidence，存入bgr_confidence:
        bgr_confidence = [0, 0, 0]
        for i in range(3):
            res_temp = cv2.matchTemplate(src_bgr[i], sch_bgr[i], cv2.TM_CCOEFF_NORMED)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res_temp)
            bgr_confidence[i] = max_val
        return min(bgr_confidence)
